import os
import uuid
import tempfile
from PIL import Image
from fastapi import APIRouter, UploadFile, File, HTTPException, Query
from typing import Optional

from py_files.image2pdfself import process_image_to_pdf, batch_process_images_to_pdf, is_image_file, ensure_directory_exists
from py_files.pdf2md import imagePdf2md

# 创建路由器
router = APIRouter(
    prefix="/image2md",
    tags=["图片转Markdown"],
    responses={404: {"description": "Not found"}},
)

# 默认目录
DEFAULT_INPUT_DIR = "input_files"
DEFAULT_OUTPUT_DIR = "output_files"

def process_image_to_md(image_path, output_dir="output_files"):
    """
    将图片转换为Markdown文件，通过先将图片转为PDF，再使用PDF提取器处理
    
    Args:
        image_path: 图片文件路径
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的Markdown文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 检查文件是否为图片
    if not is_image_file(image_path):
        raise ValueError(f"文件不是有效的图片: {image_path}")
    
    # 生成唯一的输出文件名（用于最终Markdown）
    image_basename = os.path.basename(image_path)
    base_name = os.path.splitext(image_basename)[0]
    unique_id = uuid.uuid4().hex[:8]
    output_filename = f"output_image_{base_name}_{unique_id}.md"
    markdown_path = os.path.join(output_dir, output_filename)
    
    try:
        # 首先将图片转换为临时PDF文件
        temp_dir = tempfile.mkdtemp()
        pdf_path = process_image_to_pdf(image_path, output_dir=temp_dir)
        
        # 使用PDF提取器将PDF转换为Markdown
        imagePdf2md(pdf_path, markdown_path)
        
        # 清理临时文件
        if os.path.exists(pdf_path):
            os.remove(pdf_path)
        if os.path.exists(temp_dir):
            os.rmdir(temp_dir)
        
        return markdown_path
    except Exception as e:
        # 如果处理出错，清理临时资源
        raise ValueError(f"处理图片时出错: {str(e)}")

def batch_process_images_to_md(image_paths, output_dir="output_files"):
    """
    批量将多张图片转换为单个Markdown文件，通过先将图片转为PDF，再使用PDF提取器处理
    
    Args:
        image_paths: 图片文件路径列表
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的Markdown文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 获取有效的图片文件
    valid_image_paths = []
    for image_path in image_paths:
        if is_image_file(image_path):
            valid_image_paths.append(image_path)
    
    if not valid_image_paths:
        raise ValueError("没有有效的图片文件可以转换")
    
    # 生成唯一的输出文件名
    unique_id = uuid.uuid4().hex[:8]
    output_filename = f"output_batch_images_{unique_id}.md"
    markdown_path = os.path.join(output_dir, output_filename)
    
    try:
        # 首先将多张图片合并为临时PDF文件
        temp_dir = tempfile.mkdtemp()
        pdf_path = batch_process_images_to_pdf(valid_image_paths, output_dir=temp_dir)
        
        # 使用PDF提取器将PDF转换为Markdown
        imagePdf2md(pdf_path, markdown_path)
        
        # 清理临时文件
        if os.path.exists(pdf_path):
            os.remove(pdf_path)
        if os.path.exists(temp_dir):
            os.rmdir(temp_dir)
        
        return markdown_path
    except Exception as e:
        # 如果处理出错，清理临时资源
        raise ValueError(f"处理图片时出错: {str(e)}")

# 添加路由器端点
@router.post("/convert")
async def convert_image_to_md(
    file: UploadFile = File(...),
    input_dir: Optional[str] = Query(DEFAULT_INPUT_DIR, description="输入文件目录"),
    output_dir: Optional[str] = Query(DEFAULT_OUTPUT_DIR, description="输出文件目录")
):
    """
    将图片文件转换为Markdown（先转为PDF，再使用高级PDF提取器处理）
    
    - **file**: 要上传的图片文件
    - **input_dir**: 输入文件保存目录，默认为'input_files'
    - **output_dir**: 输出文件保存目录，默认为'output_files'
    
    注意：此接口将图片先通过img2pdf无损转换为PDF，然后使用PDF提取器提取内容，
    可以保留表格结构、公式等复杂内容，提供更高质量的Markdown转换结果。
    """
    # 处理默认目录
    if not input_dir or input_dir.strip() == "":
        input_dir = DEFAULT_INPUT_DIR
    
    if not output_dir or output_dir.strip() == "":
        output_dir = DEFAULT_OUTPUT_DIR
    
    # 确保目录存在
    ensure_directory_exists(input_dir)
    ensure_directory_exists(output_dir)
    
    # 检查文件类型（通过扩展名简单判断，实际处理时会进一步验证）
    valid_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.tiff', '.tif']
    file_ext = os.path.splitext(file.filename.lower())[1]
    if file_ext not in valid_extensions:
        raise HTTPException(status_code=400, detail=f"只接受以下格式的图片文件: {', '.join(valid_extensions)}")
    
    # 保存上传的文件
    file_path = os.path.join(input_dir, file.filename)
    with open(file_path, "wb") as buffer:
        buffer.write(await file.read())
    
    try:
        # 处理图片文件
        output_path = process_image_to_md(file_path, output_dir=output_dir)
        
        # 返回文件下载URL
        file_name = os.path.basename(output_path)
        download_url = f"/outputs/{file_name}"
        
        return {
            "success": True,
            "message": "图片转换为Markdown成功（通过PDF提取）",
            "file_path": output_path,
            "download_url": download_url
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文件时出错: {str(e)}")
